# from paddlex import create_model
# model = create_model(model_name="PP-LCNet_x1_0_table_cls")
# output = model.predict("image.png", batch_size=1)
# for res in output:
#     res.print(json_format=False)
#     res.save_to_json("output/res.json")
import copy
# from paddlex import create_model
# model = create_model(model_name="SLANeXt_wireless")
# output = model.predict(input="image.png", batch_size=1)
# for res in output:
#     res.print(json_format=False)
#     res.save_to_json("output/res.json")
import time
import os
from pathlib import Path

os.environ["OMP_NUM_THREADS"] = "1"
os.environ["KMP_AFFINITY"] = "none"
import yaml
with open('table_recognition_v2_llc.yaml', 'r') as f:
    yaml_config = yaml.safe_load(f)


from paddlex import create_pipeline

pipeline = create_pipeline(pipeline="table_recognition_v2_llc.yaml")
# pipeline = create_pipeline(pipeline="table_recognition_v2.yaml")

output = pipeline.predict(
    input="3a417bf2cc6a33e21bf2cd995da6898bd1de20d6a35b24157b27a77244616b1d.jpg",
    use_doc_orientation_classify=False,
    use_doc_unwarping=False,
    # use_e2e_wireless_table_rec_model=True
)

def s_img(res):
    from PIL import Image, ImageDraw
    table_cell_img = Image.fromarray(
        copy.deepcopy(res["doc_preprocessor_res"]["output_img"])
    )
    table_draw = ImageDraw.Draw(table_cell_img)
    rectangle_color = (255, 0, 0)
    for sno in range(len(res["table_res_list"])):
        table_res = res["table_res_list"][sno]
        cell_box_list = table_res["cell_box_list"]
        for box in cell_box_list:
            x1, y1, x2, y2 = [int(pos) for pos in box]
            if x1 > 15:
                continue
            table_draw.rectangle(
                [x1, y1, x2, y2], outline=rectangle_color, width=2
            )
            time.sleep(1)
            res._img_writer.write(f"/home/fengjie/doc-parser/MinerU/src/paddlex_table/output/llc/s_cell{sno}.jpg", table_cell_img)

    img = table_cell_img



    res._img_writer.write("/home/fengjie/doc-parser/MinerU/src/paddlex_table/output/llc/s_cell.jpg", img)
for res in output:
    res.print()
    s_img(res)
    # res.save_to_img("output/llc/")
    # res.save_to_xlsx("output/llc/")
    # res.save_to_html("output/llc/")
    # res.save_to_json("output/llc/")